import serial
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import time
import os

# 修改串口设置提示，便于用户操作
def setup_serial(port=None, baudrate=115200):
    if port is None:
        port = input("请输入串口端口（例如COM3）: ")
    
    try:
        ser = serial.Serial(port, baudrate, timeout=1)
        print(f"串口连接成功: {port}")
        return ser
    except Exception as e:
        print(f"串口连接失败: {e}")
        print("请检查：\n1. 开发板是否已上电\n2. 串口驱动是否安装正确\n3. 端口号是否正确")
        return None

# 加载MNIST图像
def load_mnist_image(file_path):
    image = Image.open(file_path).convert('L')  # 转换为灰度图
    image = image.resize((28, 28))  # 确保尺寸为28x28
    img_array = np.array(image)
    # 归一化并反转颜色
    img_array = 255 - img_array
    img_array = img_array / 255.0
    return img_array

# 将图像数据发送到Zynq
def send_image_to_zynq(serial_port, image):
    # 将float转为byte
    image_bytes = (image * 255).astype(np.uint8).flatten().tobytes()
    
    # 发送开始标记
    serial_port.write(b'START\n')
    time.sleep(0.1)
    
    # 发送数据
    serial_port.write(image_bytes)
    print(f"已发送 {len(image_bytes)} 字节的图像数据")
    
    # 发送结束标记
    time.sleep(0.1)
    serial_port.write(b'END\n')

# 接收识别结果
def receive_result(serial_port):
    result = None
    while True:
        line = serial_port.readline().decode('utf-8').strip()
        if line.startswith("RESULT:"):
            result = int(line.split(':')[1])
            break
        elif line:
            print(f"接收到: {line}")
    return result

# 显示图像和结果
def display_result(image, result):
    plt.figure(figsize=(5, 5))
    plt.imshow(image, cmap='gray')
    plt.title(f"识别结果: {result}")
    plt.axis('off')
    plt.show()

def main():
    print("===== MNIST手写数字识别系统客户端 =====")
    print("请确保Zynq开发板已连接并且程序已运行")
    
    # 配置串口
    ser = setup_serial()
    if not ser:
        return
    
    while True:
        # 加载测试图像
        test_image_path = input("\n请输入MNIST测试图像路径（输入'q'退出）: ")
        if test_image_path.lower() == 'q':
            break
            
        if not os.path.exists(test_image_path):
            print("文件不存在！请检查路径")
            continue
        
        try:
            image = load_mnist_image(test_image_path)
            
            # 发送图像到Zynq
            print("发送图像到Zynq...")
            send_image_to_zynq(ser, image)
            
            # 等待并接收结果
            print("等待识别结果...")
            result = receive_result(ser)
            
            if result is not None:
                print(f"识别结果: {result}")
                display_result(image, result)
            else:
                print("未收到有效结果，请检查开发板连接")
        except Exception as e:
            print(f"处理过程中出错: {e}")
    
    # 关闭串口
    print("正在关闭连接...")
    ser.close()
    print("程序已退出")

if __name__ == "__main__":
    main()
